Add FP8×INT4 rowwise GEMM for ROCm (f8i4bf16_rowwise) - #495

This PR is a copy of apicciau#3

Summary

Adds a Triton implementation of f8i4bf16_rowwise (FP8 activations × INT4 weights → BF16 output) for AMD MI300X/MI350X. The kernel reuses the unified _bf16i4_rowwise_kernel from the BF16×INT4 path, gated on a compile-time HAS_X_SCALE constexpr — no kernel fork.

Context

• Hardware: AMD Instinct MI300X (gfx942), MI350X (gfx950, CDNA4)
• Depends on Add Meta license header to source files #1 (rocm/bf16i4-shuffled-grouped) — this PR stacks on top of that branch
• Upstream PR: Add Triton FP8×INT4 rowwise GEMM for ROCm gfx942/gfx950 #479

Implementation

Shared kernel with HAS_X_SCALE constexpr. The existing _bf16i4_rowwise_kernel gains two new parameters: x_scale_ptr (per-row FP8 activation dequant scale) and HAS_X_SCALE: tl.constexpr = False. When HAS_X_SCALE=False (BF16 path), the kernel behaves identically to before — the default value preserves backward compatibility and the dummy pointer is never dereferenced.

Deferred BF16 upcast. The .to(tl.bfloat16) cast on activation loads is moved from the load site to a single centralised location after the four masking branches. This is required because FP8 inputs return a native FP8 tensor type from tl.load that must be upcast explicitly — but for BF16 inputs the cast is a no-op, so both paths share the same code.

Two-dot path for FP8. When HAS_X_SCALE=True, the kernel issues two separate tl.dot calls (one for even-K, one for odd-K) instead of fusing via tl.cat. The tl.cat fusion concatenates tiles in LDS before the dot, which causes bank conflicts that are a net loss for FP8 activations on MI300X. Two separate dots avoid the LDS staging at the cost of two MFMA instructions — a net win for this dtype. The BF16 path retains the fused tl.cat dot (controlled by FUSE_DOT).

Per-row activation scale epilogue. After the K-loop, the float32 accumulator is multiplied element-wise by x_scale[m] before storing to the workspace. This fuses the FP8 dequantisation scale into the GEMM epilogue rather than requiring a separate pointwise kernel.

FP8 autotuner pruning. Four additional prune rules fire only when HAS_X_SCALE=True: skip BLOCK_N < 128 for large N, skip BLOCK_K > 64, skip SPLIT_K == 2, fix GROUP_SIZE_M = 4. These were determined by profiling winning configs on MI350X and reduce the FP8 autotuning search space by ~5×.

Thin Python wrapper. f8i4bf16_rowwise_gemm.py byte-splits FP8 activations into even/odd K columns, reinterprets them as FP8 TensorWrappers via reinterpret_fp8_type, and calls the shared kernel with HAS_X_SCALE=True. The wrapper also coerces scale/zp tensors to float32. Registered as mslk::f8i4bf16_rowwise on ROCm via torch.library.impl.

Testing

All 9 existing BF16×INT4 accuracy tests pass (no regressions):

9 passed in 74s

FP8×INT4 accuracy against dequantised reference (deq(FP8_x) @ deq(INT4_w)):

FP8×INT4 dispatch test (torch.ops.mslk.f8i4bf16_rowwise vs direct call): bitwise identical on all shapes.

Performance

Hardware: AMD Instinct MI350X (gfx950), ROCm 7.1.1, Triton 3.5.1+rocm7.1.1.
triton.testing.do_bench, warmup=25, rep=200.
Peak: 4600 TFLOPS FP8, 2300 TFLOPS BF16, 8000 GB/s HBM.

f8i4bf16_rowwise vs bf16i4bf16_rowwise — decode shapes (memory-bandwidth bound)

Model-attributed shapes. FP8 and BF16 paths are both kernel-launch-latency bound at small M; performance converges.

f8i4bf16_rowwise vs bf16i4bf16_rowwise — prefill shapes (compute bound)

FP8 MFMA has 2× the throughput of BF16 MFMA. The FP8 path is 5–8% faster at large M.